package com.sprint.questai.module.prompt;

import com.sprint.questai.model.entity.UserCaseInfo;
import com.sprint.questai.model.enums.ChatModelEnums;
import com.sprint.questai.service.Lawyer;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.input.Prompt;
import dev.langchain4j.model.input.PromptTemplate;
import dev.langchain4j.service.AiServices;

import java.util.HashMap;
import java.util.Map;

/**
 * @Author: mayixiang
 * @Date: 2024/6/22 15:15
 * @Description: 提取信息
 */
public class ExtractPrompt {

    static String question="我喝了酒开车撞了人，被抓获了，会怎么判刑";
    static String outputJson="[!!important]要求:结果返回json,不要加上Markdown标记,禁止使用markdown输出!该规则必须满足";
    public static String getPrompt(String info){
        PromptTemplate template= PromptTemplate.from("你要分析旧的信息并提取新的信息,以json的形式返回,注意只需要返回json,不要加上Markdown标记,禁止使用markdown输出" +
                "下面我会给你用户提供的最新信息，示例的更新信息的形式如下：我当时把对方打到骨折了,然后想要逃跑的时候被警察带走了" +
                "输出示例如下:{对方伤势:手臂骨折,是否自首:否},注意这只是一个参考示例,你要理解对方写的事实并通过json展示,注意输出为json格式,不要使用markdown" +
                "下面是用户提供的最新信息：{{info}}");
        Map<String ,Object> map=new HashMap<>();
        map.put("info",info);
        Prompt prompt = template.apply(map);
        return prompt.text();
    }
    public static UserCaseInfo updateUserInfo(UserCaseInfo userInfo,String info){
        ChatLanguageModel model = ChatModelEnums.glm3.getValue();
        Lawyer lawyer = AiServices.builder(Lawyer.class)
                .chatLanguageModel(model)
                .build();
        if(userInfo==null){
            userInfo=new UserCaseInfo();
        }
        String submittedPrompt="目前用户已经提供的信息如下:"+userInfo;
        String newPrompt=getPrompt(info);
        UserCaseInfo keyValues = lawyer.getKeyValues(outputJson+submittedPrompt + newPrompt);
        System.out.println(keyValues);
        return keyValues;
    }
    public static void main(String[] args) {
        Map<String,String> map=new HashMap<>();
        map.put("是否有驾驶证","是");
        UserCaseInfo userInfo=new UserCaseInfo();
        userInfo.setUserCaseInfo(map);
        String info="/信息 是否醉驾：是；酒精浓度：120";
        updateUserInfo(userInfo,info);
    }
}
